import argparse
import json
import config.config_ml as config
from src.utils.save import save_metrics, save_hyperparameters
from src.utils.get import get_df, get_model, get_output_dir, get_split_gen
from src.utils.visualize import plot_confusion_matrix, plot_10fold_confusion_matrix
from src.training.training_ml import train

# 解析参数
parser = argparse.ArgumentParser()
parser.add_argument("--model", type=str, default="rf",
                    choices=["rf", "xgb", "svm", "linear", "gbr"])
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--experiment", type=str, default="classification",
                    choices=["classification", "regression"])
parser.add_argument("--tune", action="store_true", help="是否进行超参数调节")
parser.add_argument("--cv", action="store_true", help="是否使用十折交叉验证")
parser.add_argument("--params", action="store_true", help="是否使用预训练参数")
parser.add_argument("--sensitive", action="store_true", help="是否使用敏感性分析数据")
args = parser.parse_args()

# === 读取数据 ===
df = get_df(args.experiment, args.sensitive)
print("df.shape:", df.shape)

# === 选择模型 ===
ModelClass = get_model(args.model, args.experiment)

# === 加载最佳参数 ===
if args.params:
    with open(f"ml_params/{args.experiment}/{args.model}.json", "r") as f:
        best_params = json.load(f)
    print(f"加载最佳参数: {best_params}")
else:
    best_params = None

# === 输出目录 ===
output_dir, csv_path = get_output_dir(
    experiment=args.experiment,
    model=args.model,
    cv=args.cv,
    sensitive=args.sensitive,
)

# === 保存超参数信息（JSON 格式）===
save_hyperparameters(args, config, output_dir, pipeline="ml", best_params=best_params)
print(f"超参数已保存到 {output_dir}/hyperparameters.json")

# === 数据切分器（统一）===
split_gen = get_split_gen(df, args.experiment, sensitive=args.sensitive, seed=args.seed, cv=args.cv)

# === 训练和评估 ===
all_metrics = []

for fold, (train_idx, test_idx) in enumerate(split_gen):
    train_df, test_df = df.iloc[train_idx], df.iloc[test_idx]
    print(f"Fold {fold+1} - Train size: {len(train_df)}, Test size: {len(test_df)}")

    model = train(train_df, ModelClass, best_params, args, config)
    metrics = model.evaluate(test_df)
    all_metrics.append(metrics)

    # 保存 csv
    if args.experiment == "classification":
        exclude_keys = ['confusion_matrix', 'classification_report', 'best_params', 'params']
        filtered_metrics = {k: v for k, v in metrics.items() if k not in exclude_keys}
        save_metrics(filtered_metrics, csv_path)
    else:
        save_metrics(metrics, csv_path)

    # 打印报告
    if args.experiment == "classification":
        print(metrics['confusion_matrix'])
        print(metrics['classification_report'])
    if not args.params and args.tune:
        print(metrics['params'])

# === 可视化/保存 ===
if args.cv:
    print(f"10fold指标已保存到 {csv_path}")
    if args.experiment == "classification":
        plot_10fold_confusion_matrix(all_metrics, f"{output_dir}/10fold_confusion_matrix.png")
else:
    print(f"平均指标已保存到 {csv_path}")
    if args.experiment == "classification":
        plot_confusion_matrix(all_metrics[0]['confusion_matrix'], f"{output_dir}/confusion_matrix.png")
    # 保存模型（仅非CV模式）
    model.save(f"{output_dir}/best_model_{args.model}.pkl")
    print(f"模型已保存到 {output_dir}/best_model_{args.model}.pkl")
